'''
data_dzh                大智慧自定义数据文本文
StockDataFetcher        获取沪深股票抓取
stock_zh_index_daily    指数抓取

MyStockData             股票数据源类
MySignalData            自定义数据源类
'''
import pandas as pd
from datetime import datetime
import akshare as ak
import backtrader as bt  

#===================ETF数据类==============================
class etf_daily:
    """
    获取股票数据，复权设置
    # 示例使用
        symbol = "510050"
        adjust = "hfq"

        etf = etf_daily(symbol,  adjust)
        data = etf.get_data()
        print(data)
    
    """
    def __init__(self, symbol, adjust):
        self.symbol = symbol
        self.adjust = adjust
        self.data = None
        
    def get_data(self):
        '''获取股票后复权数据'''
        fund_etf_hist_em_df = ak.fund_etf_hist_em(symbol=self.symbol, 
                                                  period="daily", 
                                                  adjust=self.adjust)
        selected_columns = ['日期', '开盘', '收盘', '最高', '最低', '成交量']
        etf_df = fund_etf_hist_em_df[selected_columns]
        
        # 处理字段命名
        etf_df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
        # 确保 'date' 列是 datetime 类型
        pd.options.mode.chained_assignment = None  # 关闭警告
        etf_df['date'] = pd.to_datetime(etf_df['date'])
        # 将date设为索引
        etf_df.index = pd.to_datetime(etf_df['date'])
        
        self.data = etf_df
        return self.data
# ==================大智慧自定义数据文本文==================
class data_dzh: 
    def __init__(self, file_path):
        self.file_path = file_path
        self.data = None

    def read_and_convert_txt_to_dataframe(self):
        """
        读取大智慧自定义数据文本文件并将其转换为 DataFrame，同时将 'Date' 列转换为日期格式。

        返回:
        pd.DataFrame: 转换后的 DataFrame。
        """
        # 读取文本文件
        self.data = pd.read_csv(self.file_path, delimiter='\t', header=None, names=['Date', 'Value'])
        # 将 'Date' 列转换为日期格式
        self.data['Date'] = pd.to_datetime(self.data['Date'])
        self.data.columns = ['date', 'value',]
        # 时间索引
        self.data.index = pd.to_datetime(self.data['date'])
        
        return self.data

# ==================指数抓取==================
class stock_zh_index_daily: # akshare 指数
    def __init__(self, symbol):
        self.symbol = symbol
        self.data_stock = None

    def get_data(self):
        """
        处理数据：提取需要的列、改变列名称、转换日期列类型、设置时间索引
                        date      open      high       low     close       volume
        date
        1990-12-19 1990-12-19    96.050    99.980    95.790    99.980       126000
        """
        self.data_stock = ak.stock_zh_index_daily(symbol=self.symbol)
        if self.data_stock is None:
            raise ValueError("确认股指数据 ak.stock_zh_index_daily(symbol=self.symbol).")

        # 日期列转换时间类型
        self.data_stock['date'] = pd.to_datetime(self.data_stock['date'])

        # 设置时间索引
        self.data_stock.index = pd.to_datetime(self.data_stock['date'])

        return self.data_stock

# ==================获取沪深股票抓取==================    
class StockDataFetcher:
    """获取股票数据，复权设置"""
    def __init__(self, symbol, adjust="hfq"):
        self.symbol = symbol
        self.adjust = adjust
        self.data = None

    def get_data(self):
        '''获取股票后复权数据'''
        stock_df = ak.stock_zh_a_hist(symbol=self.symbol, adjust=self.adjust)
        selected_columns = ['日期', '开盘', '收盘', '最高', '最低', '成交量']
        stock_df = stock_df[selected_columns]
        
        # 处理字段命名
        stock_df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
        # 确保 'date' 列是 datetime 类型
        stock_df['date'] = pd.to_datetime(stock_df['date'])
        # 将date设为索引
        stock_df.index = pd.to_datetime(stock_df['date'])
        
        self.data = stock_df
        return self.data

# ==================自定义数据源类==================
# 这两个类都继承自 bt.feeds.PandasData，意味着它们能够处理Pandas DataFrame格式的数据。
# 在策略中，您可以通过 self.datas[0] 访问股票数据，通过 self.datas[1] 访问信号数据。
class MyStockData(bt.feeds.PandasData):  # MyStockData 类定义了如何从Pandas DataFrame中提取股票数据。
    # params 元组指定了日期、开盘价、最高价、最低价、收盘价和交易量（如果有的话）在DataFrame中的列索引。
    params = (  
        ('datetime', 0),  # 日期在DataFrame中的列索引 
        ('open', 1),      # 开盘价在DataFrame中的列索引  
        ('close', 2),     # 收盘价在DataFrame中的列索引  
        ('high', 3),      # 最高价在DataFrame中的列索引 
        ('low', 4),       # 最低价在DataFrame中的列索引  
        ('volume', 5),    # 交易量在DataFrame中的列索引（如果有的话）    
    )  

# 自定义数据源，这里仅展示框架，具体实现需要根据你的CSV文件格式进行调整  
class MySignalData(bt.feeds.PandasData):  # MySignalData 类用于处理信号数据。
    lines = ('value',)  # lines = ('signal',) 是一个元组的示例，
    # 指定了日期和信号在DataFrame中的列索引。
    params = (  
        ('datetime', 0),  # 日期在DataFrame中的列索引
        ('value', 1), # 假设'value'是某种信号，这里重命名以避免混淆  

    ) 
    
# ====================================================================================
if __name__ == '__main__':
    # 示例调用
    data_M5M60_file_path = '法宝回测/dzh自定义/M5M60.txt'  # 替换为你的文件路径
    data_M5M60 = data_dzh(data_M5M60_file_path)
    data_M5M60 = data_M5M60.read_and_convert_txt_to_dataframe()
    print(data_M5M60.tail(1))
    
    data_stock = stock_zh_index_daily(symbol="sh000001")
    data_stock = data_stock.get_data()
    print(data_stock.tail(1))
    
    bwwmacd_file_path = '法宝回测/dzh自定义/macd.txt'  # 替换为你的文件路径
    data_bwwmacd = data_dzh(bwwmacd_file_path)
    data_bwwmacd = data_bwwmacd.read_and_convert_txt_to_dataframe()
    print(data_bwwmacd.tail(1))

    stock_fetcher = StockDataFetcher(symbol="600600")
    data_stock = stock_fetcher.get_data()
    print(data_stock.tail())
# # =============添加股票数据源  ===================================
# stock_fetcher = StockDataFetcher(symbol=symbol) # 抓取股票行情类实例化
# data_stock = stock_fetcher.get_data()           # 对象的获取数据方法
# data = MyStockData(dataname=data_stock, 
#                   fromdate=start_date, 
#                   todate=end_date) # 实例数据源 
# cerebro.adddata(data, name=symbol)  # Add the data feed